package optimization.aco;

import java.io.File;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Random;

import optimization.AbstractOptimizationAlgorithm;
import optimization.aco.config.AntColonyOptimizationConfig;
import task.AbstractAutomatonTask;
import automaton.AutomatonMetaData;
import automaton.FullNfaFactory;
import automaton.efsm.DefiniteEFSM;
import automaton.efsm.NondefiniteEFSM;

/**
 * 
 * @author Daniil Chivilikhin
 * 
 * This class implements an ant colony optimization algorithm
 * for building an optimal automata for a given fitness function.
 */
public class AntColonyOptimizationAlgorithm extends AbstractOptimizationAlgorithm {
	private Random random = new Random();
	private NondefiniteEFSM nfa;
	private ArrayList<AutomatonMetaData> bestAutomatasHistory;
	private double evaporationRate;
	private double minPheromoneValueCoefficient;
	private int numberOfStates;
	private int stagnationNumber;
	
	public AntColonyOptimizationAlgorithm(AbstractAutomatonTask task, AntColonyOptimizationConfig config) {
		this.task = task;
		bestAutomatasHistory = new ArrayList<AutomatonMetaData>();
		numberOfStates = config.numberOfNFAStates();
		stagnationNumber = config.stagnationNumber();
		minPheromoneValueCoefficient = config.minPheromoneValueCoefficient();
		evaporationRate = config.evaporationRate();
	}

	/**
	 * Chooses the most valuable transition from a list of transitions
	 * using the "" method.
	 * @param transitions a list of transitions
	 * @return the chosen transition
	 */
	private NondefiniteEFSM.Transition rouletteSelection(
			ArrayList<NondefiniteEFSM.Transition> transitions) {
		Collections.sort(transitions);
		int size = transitions.size();

		double weight[] = new double[size];
		weight[0] = transitions.get(0).getPheromone();

		for (int i = 1; i < size; i++) {
			weight[i] = weight[i - 1] + transitions.get(i).getPheromone();
		}

		double p = weight[size - 1] * random.nextDouble();
		int j = 0;
		while (p > weight[j]) {
			j++;
		}

		return transitions.get(j);
	}

	/**
	 * Launching the ants to run on the NFA, thus building the dfa
	 * @return the resulting dfa
	 */
	private DefiniteEFSM runAntThroughNFA() {
		DefiniteEFSM.Transition[][] dfaTransitions = new DefiniteEFSM.Transition[nfa
				.getNumberOfStates()][nfa.getNumberOfExternalInfluences()];

		for (int i = 0; i < nfa.getNumberOfStates(); i++) {
			int currentState = i;

			NondefiniteEFSM.Transition transitions[] = new NondefiniteEFSM.Transition[nfa
					.getNumberOfExternalInfluences()];
			for (int j = 0; j < transitions.length; j++) {
				transitions[j] = rouletteSelection(nfa.transitions[currentState][j]);
				if (transitions[j].getEndState() == -1) {
					dfaTransitions[currentState][j] = new DefiniteEFSM.Transition(transitions[j]);
					continue;
				}
				dfaTransitions[currentState][j] = new DefiniteEFSM.Transition(
						transitions[j]);
			}
		}
		return new DefiniteEFSM(nfa.getNumberOfStates(),
				nfa.getNumberOfExternalInfluences(), dfaTransitions);
	}

	/**
	 * Imprinting the pheromone on the transitions of NFA corresponding to the transitions of a DFA
	 * @param dfaMetaData a DFAMetaData object 
	 * @param fitness value of the fitness function to be imprinted
	 */
	private void imprintPheromone(AutomatonMetaData dfaMetaData) {
		for (NondefiniteEFSM.Transition tr : dfaMetaData.getVisitedTransitions()) {
			for (int k = 0; k < nfa.transitions[tr.getStartState()][tr.getSymbol()].size(); k++) {
				NondefiniteEFSM.Transition tmp = nfa.transitions[tr.getStartState()][tr.getSymbol()].get(k);
				if (tmp.equals(tr)) {
					double curP = tmp.getPheromone();
					curP += dfaMetaData.getFitness() / task.getDesiredFitness();

					if (curP < NondefiniteEFSM.MIN_PHEROMONE_VALUE) {
						curP = NondefiniteEFSM.MIN_PHEROMONE_VALUE;
					}

					tmp.setPheromone(curP);
					nfa.transitions[tr.getStartState()][tr.getSymbol()].set(k, tmp);
				}
			}
		}
	}

	/**
	 * Evaporation of the pheromone from the transitions of the NFA
	 * @return maximum value of pheromone in the NFA after evaporation
	 */
	private double evaporatePheromone() {
		double maxPheromoneValue = 0;
		for (int i = 0; i < nfa.getNumberOfStates(); i++) {
			for (int j = 0; j < nfa.getNumberOfExternalInfluences(); j++) {
				for (int k = 0; k < nfa.transitions[i][j].size(); k++) {
					NondefiniteEFSM.Transition transition = nfa.transitions[i][j].get(k);
					double currentPheromone = transition.getPheromone();
					currentPheromone = (1.0 - evaporationRate) * currentPheromone;

					if (currentPheromone < NondefiniteEFSM.MIN_PHEROMONE_VALUE) {
						currentPheromone = NondefiniteEFSM.MIN_PHEROMONE_VALUE;
					}
					if (currentPheromone > maxPheromoneValue) {
						maxPheromoneValue = currentPheromone;
					}
					transition.setPheromone(currentPheromone);
					nfa.transitions[i][j].set(k, transition);
				}
			}
		}
		return maxPheromoneValue;
	}

	/**
	 * Launching the ant colony optimization algorithm.
	 * @return built finite automaton with metadata.
	 */
	@Override
	public AutomatonMetaData run() {
		long start = System.currentTimeMillis();
		nfa = FullNfaFactory.getFullAutomata(numberOfStates, task.getEvents().size(), task.getActions());
		bestAutomatasHistory = new ArrayList<AutomatonMetaData>();
		double maxPheromoneValue = NondefiniteEFSM.MIN_PHEROMONE_VALUE;
		String resultDirName = "er=" + evaporationRate;	
		File resultDir = new File(resultDirName);
		resultDir.mkdir();

		int stepCounter = 0;
		double bestFitness = -1000;
		AutomatonMetaData bestDFA = null;
		int lastBestFitnessOccurence = 0;
		while (true) {
			stepCounter++;

			if (stepCounter - lastBestFitnessOccurence == stagnationNumber) {
				System.out.println("Stagnation. Restarting....");
				lastBestFitnessOccurence = stepCounter;
				nfa = FullNfaFactory.getFullAutomata(numberOfStates, task.getEvents().size(), task.getActions());
				bestFitness = -1000;
				bestDFA = null;
			}
			
			DefiniteEFSM dfa = runAntThroughNFA();
			AutomatonMetaData dfaMetaData = task.getFitness(dfa);
			
			dfa = dfaMetaData.getDFA();
			double fitness = dfaMetaData.getFitness();
			
			if (fitness - bestFitness > 1e-10) {
				lastBestFitnessOccurence = stepCounter;
				bestDFA = dfaMetaData;
				bestFitness = fitness;
				bestDFA.setNumberOfSteps(stepCounter);
				bestAutomatasHistory.add(bestDFA);
				
				if (bestFitness > task.getMinimumFitness()) {
					bestDFA.printTransitionDiagram(resultDirName);
					bestDFA.printToGraphViz(resultDirName, task.getEvents());
				}
			}
			
			maxPheromoneValue = evaporatePheromone();
			imprintPheromone(dfaMetaData);
			NondefiniteEFSM.MIN_PHEROMONE_VALUE = minPheromoneValueCoefficient * maxPheromoneValue;
			imprintPheromone(bestDFA);

//			if (stepCounter % 100 == 0) {
//				System.out.println(stepCounter + ": current = " + fitness
//						+ "; best = " + bestFitness + "; max = "
//						+ maxPheromoneValue);
//			}

			if (bestFitness >= task.getDesiredFitness()) {
				bestDFA.setNumberOfSteps(stepCounter);
				bestDFA.setTime((System.currentTimeMillis() - start) / 1e3);
				bestDFA.setFitnessHistory(bestAutomatasHistory);
				return bestDFA;
			}
		}
	}
}